import datetime
from pathlib import Path

import torch
import torchvision
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import torch.nn.functional as F
from common import *

# 2d 卷积操作函数
from torch import nn
from torchvision import transforms


def corr2d(image, K):
    """
    :param image: 输入，shape(batch_size,H,W)
    :param K: 卷积核，shape(k_h,k_w)
    :return:
    """
    batch_size, H, W = image.shape
    k_h, k_w = K.shape
    # 经滑动卷积操作后得到的新图尺寸
    new_h = H - k_h + 1
    new_w = W - k_w + 1
    Y = torch.zeros((batch_size, new_h, new_w)).to(device)
    # 进行卷积操作，实则是对应的窗口覆盖下的矩阵对应元素值相乘
    for i in range(Y.shape[1]):
        for j in range(Y.shape[2]):
            Y[:, i, j] = (image[:, i: i + k_h, j: j + k_w] * K).sum()

    return Y


def corr2d_multi_in(X, K):
    # 输入X 维度（batch_size, C_in,H,W)
    # 卷积核K: 维度(C_in,k_h,k_w)
    # 输出：维度（batch_size,H_out,w_out)
    res = corr2d(X[:, 0, :, :], K[0, :, :])
    for i in range(1, X.shape[1]):
        res += corr2d(X[:, i, :, :], K[i, :, :])
    return res


def corr2d_multi_in_out(X, K):
    # X: shape (batch_size, C_in, H, W)
    # K: shape (C_out,C_in,h,w)
    # Y: shape(batch_size, C_out, H_out,W_out)
    return torch.stack([corr2d_multi_in(X, k) for k in K], dim=1)


class MyConv2D(torch.nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size):
        super(MyConv2D, self).__init__()
        # 初始化卷积层的2个参数：卷积核、偏差
        if isinstance(kernel_size, int):
            kernel_size = (kernel_size, kernel_size)
        self.weight = torch.nn.Parameter(torch.randn((out_channels, in_channels) + kernel_size))
        self.bias = torch.nn.Parameter(torch.randn(out_channels, 1, 1))

    def forward(self, x):
        """
        :param x: 输入图片，维度（batch_size,C_in,H,W)
        :return:
        """
        return corr2d_multi_in_out(x, self.weight) + self.bias


# 自定义卷积层
class MyConvModule(nn.Module):
    def __init__(self):
        super(MyConvModule, self).__init__()
        # 定义三层卷积
        self.conv = nn.Sequential(
            MyConv2D(in_channels=3, out_channels=32, kernel_size=3),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True)
        )
        # 输出层， 将通道数变为分类数量
        self.fc = nn.Linear(32, num_classes)

    def forward(self, X):
        out = self.conv(X)
        out = F.avg_pool2d(out, i_size-2)
        out = out.squeeze()
        out = self.fc(out)
        return out


# 训练函数
def train_epoch(net, data_loader, device, is_train):
    if is_train:
        net.train()  # 指定当前为训练模式
    else:
        net.eval()  # 指定当前为测试模式
    train_batch_num = len(data_loader)  # 记录共有多少个batch
    total_loss = 0  # 记录Loss
    correct = 0  # 记录共有多少个样本被正确分类
    sample_num = 0  # 记录样本总数

    for batch_idx, (data, target) in enumerate(data_loader):
        data = data.to(device).float()
        target = target.to(device).long()
        optimizer.zero_grad()
        output = net(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
        prediction = torch.argmax(output, 1)
        correct += (prediction == target).sum().item()
        sample_num += len(prediction)
    loss = total_loss / train_batch_num
    acc = correct / sample_num
    return loss, acc


# data_dir = "/Users/ryantong/Documents/bjtu/实验三数据集/车辆分类数据集"
data_dir = "/Users/ryantong/Documents/bjtu/实验三数据集/Alex"
ratio = 0.7
i_size = 32
transform = transforms.Compose(
    [transforms.Resize((i_size, i_size), Image.ANTIALIAS)
     , transforms.ToTensor()
     , transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
     ]
)
all_datas = torchvision.datasets.ImageFolder(data_dir, transform=transform)
classes = all_datas.classes
class_to_idx = all_datas.class_to_idx
sum_len = all_datas.__len__()
# 划分数据集
train_size = int(0.7 * len(all_datas))
test_size = len(all_datas) - train_size
train_dataset, test_dataset = torch.utils.data.random_split(all_datas, [train_size, test_size])
train_dataset = train_dataset.dataset  # 这行很重要
test_dataset = test_dataset.dataset  # 这行很重要
# for images, labels in train_dataset.imgs:
#     print(images)
#     print(labels)

num_classes = 3
epochs = 5
lr = 0.01
batch_size = 32
device = torch.device("cpu")
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)
print(len(train_loader))
net = MyConvModule().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), lr=lr)

train_loss_list = []
train_acc_list = []
test_loss_list = []
test_acc_list = []

start_time = datetime.datetime.now()
print(f"start training，start time = {start_time}，epochs={epochs}, lr={lr}, batch_size={batch_size}, len(train_loader)={len(train_loader)}")
for epoch in range(epochs):
    train_loss, train_acc = train_epoch(net, data_loader=train_loader, device=device, is_train=True)
    test_loss, test_acc = train_epoch(net, data_loader=test_loader, device=device, is_train=False)
    train_loss_list.append(train_loss)
    train_acc_list.append(train_acc)
    test_loss_list.append(test_loss)
    test_acc_list.append(test_acc)
    print(f"epoch:{epoch}\t train_loss:{train_loss:.4f} \t"
          f"train_acc:{train_acc} \t"
          f"epoch:{epoch}\t test_loss:{test_loss:.4f} \t"
          f"test_acc:{test_acc} \t"
          )
end_time = datetime.datetime.now()
time = end_time - start_time
print(f"end training，end time = {end_time}, duration time = {time}")
loss_draw(range(1, epochs + 1), train_loss_list, "epochs", "loss", range(1, epochs + 1), test_loss_list,
              ['train', 'test'])
